Day 3 | 限流器 Lab 3 實作：Token Bucket

17th鐵人賽

y___che

2025-08-24 16:27:09

113 瀏覽

分享至

前言

前兩天我們已經實作完「固定窗口」跟「滑動窗口」計數器這兩個好兄弟了，接下來要實作的 Token Bucket 以及 Leaky Bucket 從根本的設計哲學上，是完全獨立於之前實作過的時間窗口相關的算法，這兩天開始會說明與探討這兩種算法的實作方式。

實作 Lab 3：Token Bucket 限流算法

在接觸一個新概念時，先快速建立抽象的高層次理解是個不錯的第一步，概念的粒度越粗、包含越少的細節越好，最好是能更生活化，只要方向對即可，但其實錯了也很好，哈哈。

講這麼多，我想說的是時間窗口類型的限流原則是「多少時間我就做多少事，我就算有能力多做，我也不想多做，你晚點再來」，而 Token Bucket 的限流原則是「什麼時候來隨便你，反正我當下能做就做，不能做你就晚點來」，而背後的邏輯設計主要就是針對，「當下能做就做」跟「晚點來」做定義：

當下能做是什麼意思？怎樣叫能做？
晚點來是多晚？

在 Token Bucket 的規則裡，會有一個固定容量的桶子，裡面裝著很多 Token，請求進來發現桶子裡有 Token 就拿走、領著它繼續往下執行業務邏輯，發現桶裡沒東西了就只好下次請求時再來，而桶子會以固定的速率回存 Token，沒有請求時也會慢慢累積為流量高峰作儲備，等流量高峰一來，也不管多短時間做多少請求，有 Token 就拿去用，再大量的請求都不怕，只要桶裡還有東西就盡情使用，但若發現桶裡沒東西的話，請求就會被丟棄，而一旦桶子裝滿就不會再繼續填入了。這種限流規則很彈性、平滑，有多少能力就做多少事的概念，在這個概念裡有幾個重要參數需要去設置：

桶的容量
回存的速率
每次回存的數量

既然提到了這幾個參數，那我們就從 Bucket 的資料結構開始實作吧！

Bucket

在 TokenBucket 裡面我們定義了四個屬性：

capacity：桶子的容量（可以裝幾個 Token）
refillRate：Token 回填的速率
tokens：Token 數量
lastRefillTime：最近一次回填時間

並在下面定義一個建構子，在創建實例時我們需要讓 tokens 屬性預設等於 capacity 的大小，也就是剛被創建的桶子一開始就要是滿的，lastRefillTime 定義為當下時間。

@Getter
@Setter
public class TokenBucket {

    private final int capacity;
    private final int refillRate;
    private double tokens;
    private long lastRefillTime;

    public TokenBucket(int capacity, int refillRate) {
        this.capacity = capacity;
        this.refillRate = refillRate;
        this.tokens = capacity;
        this.lastRefillTime = System.currentTimeMillis();
    }
}

計數器

Token Bucket 跟時間窗口類型的計數器底層設計哲學差別很大，很難把實作的方法一一映射在一起說明，就像這個 tryConsumeToken 作為限流邏輯的主要方法，不像之前計數器是根據請求次數逐一遞增，而是剛好相反的用請求「消費」Token 的方式來更新限流額度，再搭配 Lazy Refill 方法計算出要回填的 Tokens 並存入桶中。至於這邊為何要用 synchronized 其實就可以參照前兩個限流算法來理解緣由了。

`tryConsumeToken` synchronized

Token Bucket 的消費跟回填機制，類似於時間窗口的 incrementAndGet，只是 incrementAndGet 是原子性的，等同消費跟回填一次做完，但是 Token Bucket無法，因為它不只是單純的遞增遞減的簡單運算，而是需要在消費 Token 之前計算當前回填 Token 的數量，只能用同步機制來確保原子性避免在多線程的環境下出現 Race Condition，導致數據不一致、重複回填等情況。主要還是因為 Lazy Refill 的機制，回填的時機點是在請求消費 Token 時才去計算的，當然也可以將消費跟回填機制分開，設置一個背景執行緒的定時任務來定期回填 Token，只是這個資源消耗很不必要，所以才需要用 synchronized 將其綁定成原子操作。

@Override
public boolean tryConsumeToken(String key, int capacity, int refillRate) {
    var bucket = tokenBucketStorage.computeIfAbsent(key, k -> new TokenBucket(capacity, refillRate));

    synchronized (bucket) {
        refillTokens(bucket);
        if (bucket.getTokens() >= 1) {
            bucket.setTokens(bucket.getTokens() - 1);
            return true;
        }
    }
    return false;
}

private void refillTokens(TokenBucket bucket) {
    var now = System.currentTimeMillis();
    var elapsedTime = now - bucket.getLastRefillTime();

    if (elapsedTime > 0) {
        var tokensToAdd = (elapsedTime / 1000.0) * bucket.getRefillRate();
        var newTokens = Math.min(bucket.getCapacity(), bucket.getTokens() + tokensToAdd);
        bucket.setTokens(newTokens);
        bucket.setLastRefillTime(now);
    }
}

`refillTokens` 計算回填的數量

那麼如何計算回填 Token 的數量呢？首先我們需要推算出上次回填到現在已經過多久了，然後轉換成秒數後跟回填速率相乘即可：

var elapsedTime = now - bucket.getLastRefillTime(); ：算出上次到現在過了多少時間
var tokensToAdd = (elapsedTime / 1000.0) * bucket.getRefillRate(); ：除掉 1000 轉換成秒，回填速率指的是一秒可以存幾個 Tokens，所以秒數跟速率相乘可得到要回填的 Tokens
var newTokens = Math.min(bucket.getCapacity(), bucket.getTokens() + tokensToAdd); ：用 Math.min 確保回填的數量不會超出桶子容量
存入、搞定

清除閒置許久的 Bucket

定時 remove 掉閒置過長時間的桶子：

// 清除策略
var beforeOneHour = System.currentTimeMillis() - 3600000;
tokenBucketStorage.entrySet().removeIf(entry -> {
    var bucket = entry.getValue();
    return bucket.getLastRefillTime() < beforeOneHour;
});

Limiter Strategy 實例

當我們實際用 @RateLimiter 定義 Token Bucket 算法時，limit 代表我們的桶子容量，window 是多久回填一次：

@RateLimiter(key = "token-bucket", limit = 5, window = 10, algorithm = Algorithm.TOKEN_BUCKET)
@GetMapping(value = "/token-bucket")
public ResponseEntity<String> getTokenBucketString() {
    return ResponseEntity.ok("Token Bucket limiter");
}

因此在算法實例的類別中，計算回填速率的方式就是：

桶子容量 / 回填週期（多久回填） ⇒ var refillRate = Math.max(1, rateLimiter.limit() / rateLimiter.window()); 這樣就可算出一秒要回填幾個 Tokens 了。下面使用 Math.max 是確保最少每秒回填 1 個 Token。

@RequiredArgsConstructor
@Component
public class TokenBucketLimiter implements RateLimiterStrategy {

    private final RateLimiterStorage storage;

    @Override
    public boolean isAllow(String key, RateLimiter rateLimiter) {
        var capacity = rateLimiter.limit();
        var refillRate = Math.max(1, rateLimiter.limit() / rateLimiter.window());
        return storage.tryConsumeToken(key, capacity, refillRate);
    }

    @Override
    public Algorithm getAlgorithmType() {
        return Algorithm.TOKEN_BUCKET;
    }

}